import paddle
import bisect
import numpy as np
import albumentations
from PIL import Image


class ConcatDatasetWithIndex(paddle.io.ConcatDataset):
    """Modified from original pytorch code to return dataset idx"""

    def __getitem__(self, idx):
        if idx < 0:
            if -idx > len(self):
                raise ValueError(
                    'absolute value of index should not exceed dataset length')
            idx = len(self) + idx
        dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx)
        if dataset_idx == 0:
            sample_idx = idx
        else:
            sample_idx = idx - self.cumulative_sizes[dataset_idx - 1]
        return self.datasets[dataset_idx][sample_idx], dataset_idx


class ImagePaths(paddle.io.Dataset):

    def __init__(self, paths, size=None, random_crop=False, labels=None):
        self.size = size
        self.random_crop = random_crop
        self.labels = dict() if labels is None else labels
        self.labels['file_path_'] = paths
        self._length = len(paths)
        if self.size is not None and self.size > 0:
            self.rescaler = paddle.vision.transforms.Resize(self.size)
            if not self.random_crop:
                self.cropper = paddle.vision.transforms.CenterCrop((self.size,
                    self.size))
            else:
                self.cropper = paddle.vision.transforms.RandomCrop((self.size,
                    self.size))
            self.preprocessor = paddle.vision.transforms.Compose([self.
                rescaler, self.cropper])
        else:
            self.preprocessor = lambda **kwargs: kwargs

    def __len__(self):
        return self._length

    def preprocess_image(self, image_path):
        image = Image.open(image_path)
        if not image.mode == 'RGB':
            image = image.convert('RGB')
        image = self.preprocessor(image)
        image = np.array(image)
        image = (image / 127.5 - 1.0).astype(np.float32)
        return image

    def __getitem__(self, i):
        example = dict()
        example['image'] = self.preprocess_image(self.labels['file_path_'][i])
        for k in self.labels:
            example[k] = self.labels[k][i]
        return example


class NumpyPaths(ImagePaths):

    def preprocess_image(self, image_path):
        image = np.load(image_path).squeeze(0)
        image = np.transpose(image, (1, 2, 0))
        image = Image.fromarray(image, mode='RGB')
        image = np.array(image).astype(np.uint8)
        image = self.preprocessor(image=image)['image']
        image = (image / 127.5 - 1.0).astype(np.float32)
        return image
